// SPDX-License-Identifier: GPL-2.0

#include <linux/module.h>
#include <linux/miscdevice.h>
#include <rdma/ib.h>
#include <rdma/ib_verbs.h>
#include <rdma/restrack.h>
#include <trace/hooks/rdma_pangea.h>

#include "restrack.h"
#include "ucm_priv.h"

static inline int rdma_restrack_get_qp_count_by_pid(struct ib_device *dev, pid_t pid)
{
	int count = 0;
	struct rdma_restrack_entry *res;
	unsigned long index;
	struct rdma_restrack_root *rt = &dev->res[RDMA_RESTRACK_QP];

	xa_lock(&rt->xa);
	xa_for_each(&rt->xa, index, res) {
		if (!res->user || !res->task)
			continue;

		if (res->task->tgid == pid)
			count++;
	}

	xa_unlock(&rt->xa);

	return count;
}

static inline int ib_get_qp_count_by_pid(pid_t pid)
{
	struct ib_device *dev;
	unsigned long index;
	struct xarray *device_list;
	xa_mark_t flag;
	int count = 0;

	down_read_devices_rwsem();
	device_list = obtain_device_list();
	flag = obtain_device_registered_flag();
	xa_for_each_marked(device_list, index, dev, flag) {
		count += rdma_restrack_get_qp_count_by_pid(dev, pid);
	}
	up_read_devices_rwsem();

	return count;
}

static ssize_t check_qp_write(struct file *filp, const char __user *buf,
			  size_t len, loff_t *pos)
{
	pid_t pid;
	int ret;

	if (!ib_safe_file_access(filp)) {
		ofed_log_print(PANGEA_ERR, "Process %d (%s) changed security contexts.",
			task_tgid_vnr(current), current->comm);
		return -EACCES;
	}

	if (len != sizeof(pid_t)) {
		ofed_log_print(PANGEA_ERR, "Invalid pid length(%zu).", len);
		return -EINVAL;
	}

	if (copy_from_user(&pid, buf, sizeof(pid_t))) {
		ofed_log_print(PANGEA_ERR, "Copy from user error, buf(%p), length(%zu).", buf, len);
		return -EFAULT;
	}

	ret = ib_get_qp_count_by_pid(pid);
	if (ret)
		ofed_log_print(PANGEA_INFO, "QP not cleaned, pid(%d), qp count(%d).", pid, ret);

	return ret == 0 ? (ssize_t)len : -EAGAIN;
}

static const struct file_operations check_qp_fops = {
	.owner	= THIS_MODULE,
	.open	= nonseekable_open,
	.write	= check_qp_write,
	.llseek	= no_llseek,
};

static struct miscdevice check_qp_misc = {
	.minor		= MISC_DYNAMIC_MINOR,
	.name		= "ib_ucm",
	.nodename	= "infiniband/ib_ucm",
	.mode		= 0666,
	.fops		= &check_qp_fops,
};

static int __init ib_ucm_init(void)
{
	int ret;

	ret = ib_ext_fd_init();
	if (ret) {
		ofed_log_print(PANGEA_ERR, "Init ucm fd module failed, ret(%d).", ret);
		return ret;
	}

	ret = misc_register(&check_qp_misc);
	if (ret) {
		ofed_log_print(PANGEA_ERR, "Init checkqp module failed, ret(%d).", ret);
		goto fd_exit;
	}

	return 0;
fd_exit:
	ib_ext_fd_exit();
	return ret;
}

static void __exit ib_ucm_cleanup(void)
{
	misc_deregister(&check_qp_misc);
	ib_ext_fd_exit();
}

MODULE_LICENSE("Dual BSD/GPL");
module_init(ib_ucm_init);
module_exit(ib_ucm_cleanup);

